//! WebSocket实时通信模块
//! 
//! 这个模块实现了WebSocket实时通信功能，包括：
//! - WebSocket连接管理
//! - 消息广播
//! - 房间管理
//! - 心跳检测

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    response::Response,
};
use futures_util::{sink::SinkExt, stream::StreamExt};
use std::sync::Arc;
use tokio::sync::broadcast;
use uuid::Uuid;
use std::collections::{HashMap, HashSet};

use crate::{
    models::{WebSocketMessage, WebSocketMessageType},
    errors::{AppError, AppResult},
    AppState,
};

/// WebSocket连接管理器
/// 
/// 负责管理所有WebSocket连接和消息广播
#[derive(Clone)]
pub struct WebSocketManager {
    /// 用户ID到发送器的映射
    connections: Arc<tokio::sync::RwLock<HashMap<Uuid, broadcast::Sender<String>>>>,
    /// 房间ID到用户ID集合的映射
    rooms: Arc<tokio::sync::RwLock<HashMap<Uuid, HashSet<Uuid>>>>,
    /// 用户ID到房间ID集合的映射
    user_rooms: Arc<tokio::sync::RwLock<HashMap<Uuid, HashSet<Uuid>>>>,
}

impl WebSocketManager {
    /// 创建新的WebSocket管理器
    pub fn new() -> Self {
        Self {
            connections: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            rooms: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            user_rooms: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
        }
    }

    /// 添加WebSocket连接
    /// 
    /// # 参数
    /// - `user_id`: 用户ID
    /// - `sender`: 消息发送器
    pub async fn add_connection(&self, user_id: Uuid, sender: broadcast::Sender<String>) {
        let mut connections = self.connections.write().await;
        connections.insert(user_id, sender);
        
        tracing::info!("✅ WebSocket连接已添加: {}", user_id);
    }

    /// 移除WebSocket连接
    /// 
    /// # 参数
    /// - `user_id`: 用户ID
    pub async fn remove_connection(&self, user_id: Uuid) {
        // 移除连接
        {
            let mut connections = self.connections.write().await;
            connections.remove(&user_id);
        }

        // 从所有房间中移除用户
        {
            let mut user_rooms = self.user_rooms.write().await;
            if let Some(rooms) = user_rooms.remove(&user_id) {
                let mut rooms_map = self.rooms.write().await;
                for room_id in rooms {
                    if let Some(users) = rooms_map.get_mut(&room_id) {
                        users.remove(&user_id);
                        
                        // 通知房间其他用户
                        Box::pin(self.broadcast_user_left(room_id, user_id)).await;
                    }
                }
            }
        }

        tracing::info!("❌ WebSocket连接已移除: {}", user_id);
    }

    /// 用户加入房间
    /// 
    /// # 参数
    /// - `user_id`: 用户ID
    /// - `room_id`: 房间ID
    pub async fn join_room(&self, user_id: Uuid, room_id: Uuid) {
        // 添加到房间
        {
            let mut rooms = self.rooms.write().await;
            rooms.entry(room_id).or_insert_with(HashSet::new).insert(user_id);
        }

        // 添加到用户房间列表
        {
            let mut user_rooms = self.user_rooms.write().await;
            user_rooms.entry(user_id).or_insert_with(HashSet::new).insert(room_id);
        }

        // 通知房间其他用户
        self.broadcast_user_joined(room_id, user_id).await;

        tracing::info!("🏠 用户加入房间: {} -> {}", user_id, room_id);
    }

    /// 用户离开房间
    /// 
    /// # 参数
    /// - `user_id`: 用户ID
    /// - `room_id`: 房间ID
    pub async fn leave_room(&self, user_id: Uuid, room_id: Uuid) {
        // 从房间移除
        {
            let mut rooms = self.rooms.write().await;
            if let Some(users) = rooms.get_mut(&room_id) {
                users.remove(&user_id);
            }
        }

        // 从用户房间列表移除
        {
            let mut user_rooms = self.user_rooms.write().await;
            if let Some(rooms) = user_rooms.get_mut(&user_id) {
                rooms.remove(&room_id);
            }
        }

        // 通知房间其他用户
        self.broadcast_user_left(room_id, user_id).await;

        tracing::info!("🚪 用户离开房间: {} -> {}", user_id, room_id);
    }

    /// 向房间广播消息
    /// 
    /// # 参数
    /// - `room_id`: 房间ID
    /// - `message`: 消息内容
    pub async fn broadcast_to_room(&self, room_id: Uuid, message: String) {
        let rooms = self.rooms.read().await;
        if let Some(users) = rooms.get(&room_id) {
            let connections = self.connections.read().await;
            let mut failed_users = Vec::new();

            for &user_id in users {
                if let Some(sender) = connections.get(&user_id) {
                    if sender.send(message.clone()).is_err() {
                        failed_users.push(user_id);
                    }
                }
            }

            // 清理失败的连接
            if !failed_users.is_empty() {
                drop(connections);
                drop(rooms);
                for user_id in failed_users {
                    Box::pin(self.remove_connection(user_id)).await;
                }
            }
        }
    }

    /// 向特定用户发送消息
    /// 
    /// # 参数
    /// - `user_id`: 用户ID
    /// - `message`: 消息内容
    pub async fn send_to_user(&self, user_id: Uuid, message: String) {
        let connections = self.connections.read().await;
        if let Some(sender) = connections.get(&user_id) {
            if sender.send(message).is_err() {
                drop(connections);
                    Box::pin(self.remove_connection(user_id)).await;
            }
        }
    }

    /// 获取房间内的用户列表
    /// 
    /// # 参数
    /// - `room_id`: 房间ID
    /// 
    /// # 返回值
    /// - `Vec<Uuid>`: 用户ID列表
    pub async fn get_room_users(&self, room_id: Uuid) -> Vec<Uuid> {
        let rooms = self.rooms.read().await;
        rooms.get(&room_id)
            .map(|users| users.iter().cloned().collect())
            .unwrap_or_default()
    }

    /// 获取用户所在的房间列表
    /// 
    /// # 参数
    /// - `user_id`: 用户ID
    /// 
    /// # 返回值
    /// - `Vec<Uuid>`: 房间ID列表
    pub async fn get_user_rooms(&self, user_id: Uuid) -> Vec<Uuid> {
        let user_rooms = self.user_rooms.read().await;
        user_rooms.get(&user_id)
            .map(|rooms| rooms.iter().cloned().collect())
            .unwrap_or_default()
    }

    /// 获取在线用户数量
    /// 
    /// # 返回值
    /// - `usize`: 在线用户数量
    pub async fn get_online_users_count(&self) -> usize {
        self.connections.read().await.len()
    }

    /// 广播用户加入消息
    async fn broadcast_user_joined(&self, room_id: Uuid, user_id: Uuid) {
        let message = WebSocketMessage {
            message_type: WebSocketMessageType::UserJoined,
            content: format!("用户 {} 加入了房间", user_id),
            sender_id: Some(user_id),
            room_id: Some(room_id),
            timestamp: chrono::Utc::now(),
        };

        if let Ok(message_str) = serde_json::to_string(&message) {
            self.broadcast_to_room(room_id, message_str).await;
        }
    }

    /// 广播用户离开消息
    async fn broadcast_user_left(&self, room_id: Uuid, user_id: Uuid) {
        let message = WebSocketMessage {
            message_type: WebSocketMessageType::UserLeft,
            content: format!("用户 {} 离开了房间", user_id),
            sender_id: Some(user_id),
            room_id: Some(room_id),
            timestamp: chrono::Utc::now(),
        };

        if let Ok(message_str) = serde_json::to_string(&message) {
            self.broadcast_to_room(room_id, message_str).await;
        }
    }
}

/// WebSocket处理器
/// 
/// 处理WebSocket升级和连接管理
pub async fn websocket_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
) -> Response {
    ws.on_upgrade(|socket| websocket_connection(socket, state))
}

/// WebSocket连接处理
/// 
/// # 参数
/// - `socket`: WebSocket连接
/// - `state`: 应用程序状态
async fn websocket_connection(socket: WebSocket, state: AppState) {
    let (mut sender, mut receiver) = socket.split();
    let user_id = Uuid::new_v4(); // 在实际应用中从JWT令牌中获取

    // 创建消息通道
    let (tx, mut rx) = broadcast::channel(1000);

    // 将连接添加到连接管理器
    state.websocket.add_connection(user_id, tx.clone()).await;

    // 发送任务
    let send_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            if sender.send(Message::Text(msg.into())).await.is_err() {
                break;
            }
        }
    });

    // 接收任务
    let state_clone = state.clone();
    let recv_task = tokio::spawn(async move {
        while let Some(msg) = receiver.next().await {
            match msg {
                Ok(Message::Text(text)) => {
                    if let Err(e) = handle_message(&text, user_id, &state_clone).await {
                        tracing::error!("❌ 处理WebSocket消息失败: {}", e);
                    }
                }
                Ok(Message::Close(_)) => {
                    tracing::info!("🔚 WebSocket连接关闭: {}", user_id);
                    break;
                }
                Err(e) => {
                    tracing::error!("❌ WebSocket错误: {}", e);
                    break;
                }
                _ => {}
            }
        }
    });

    // 等待任务完成
    tokio::select! {
        _ = send_task => {},
        _ = recv_task => {},
    }

    // 清理连接
    state.websocket.remove_connection(user_id).await;
}

/// 处理WebSocket消息
/// 
/// # 参数
/// - `text`: 消息文本
/// - `user_id`: 用户ID
/// - `state`: 应用程序状态
async fn handle_message(text: &str, user_id: Uuid, state: &AppState) -> AppResult<()> {
    let message: WebSocketMessage = serde_json::from_str(text)?;

    match message.message_type {
        WebSocketMessageType::Chat => {
            handle_chat_message(message, user_id, state).await?;
        }
        WebSocketMessageType::UserJoined => {
            handle_user_joined(message, user_id, state).await?;
        }
        WebSocketMessageType::UserLeft => {
            handle_user_left(message, user_id, state).await?;
        }
        WebSocketMessageType::Ping => {
            handle_ping(user_id, state).await?;
        }
        _ => {
            tracing::warn!("⚠️ 未知消息类型: {:?}", message.message_type);
        }
    }

    Ok(())
}

/// 处理聊天消息
async fn handle_chat_message(
    message: WebSocketMessage,
    user_id: Uuid,
    state: &AppState,
) -> AppResult<()> {
    let room_id = message.room_id.ok_or_else(|| AppError::validation("聊天消息必须指定房间ID"))?;

    // 验证用户是否在房间中
    let user_rooms = state.websocket.get_user_rooms(user_id).await;
    if !user_rooms.contains(&room_id) {
        return Err(AppError::validation("用户不在指定房间中"));
    }

    // 创建聊天消息
    let chat_message = WebSocketMessage {
        message_type: WebSocketMessageType::Chat,
        content: message.content,
        sender_id: Some(user_id),
        room_id: Some(room_id),
        timestamp: chrono::Utc::now(),
    };

    // 广播消息到房间
    let broadcast_message = serde_json::to_string(&chat_message)?;
    state.websocket.broadcast_to_room(room_id, broadcast_message).await;

    tracing::info!("💬 聊天消息: 用户 {} 在房间 {} 发送消息", user_id, room_id);
    Ok(())
}

/// 处理用户加入房间消息
async fn handle_user_joined(
    message: WebSocketMessage,
    user_id: Uuid,
    state: &AppState,
) -> AppResult<()> {
    let room_id = message.room_id.ok_or_else(|| AppError::validation("加入房间必须指定房间ID"))?;

    // 加入房间
    state.websocket.join_room(user_id, room_id).await;

    Ok(())
}

/// 处理用户离开房间消息
async fn handle_user_left(
    message: WebSocketMessage,
    user_id: Uuid,
    state: &AppState,
) -> AppResult<()> {
    if let Some(room_id) = message.room_id {
        state.websocket.leave_room(user_id, room_id).await;
    } else {
        // 离开所有房间
        let user_rooms = state.websocket.get_user_rooms(user_id).await;
        for room_id in user_rooms {
            state.websocket.leave_room(user_id, room_id).await;
        }
    }

    Ok(())
}

/// 处理心跳消息
async fn handle_ping(user_id: Uuid, state: &AppState) -> AppResult<()> {
    let pong_message = WebSocketMessage {
        message_type: WebSocketMessageType::Pong,
        content: "pong".to_string(),
        sender_id: None,
        room_id: None,
        timestamp: chrono::Utc::now(),
    };

    let message_str = serde_json::to_string(&pong_message)?;
    state.websocket.send_to_user(user_id, message_str).await;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_websocket_manager() {
        let manager = WebSocketManager::new();
        let user_id = Uuid::new_v4();
        let room_id = Uuid::new_v4();

        // 创建测试发送器
        let (tx, _rx) = broadcast::channel(1000);

        // 测试添加连接
        manager.add_connection(user_id, tx).await;
        assert_eq!(manager.get_online_users_count().await, 1);

        // 测试加入房间
        manager.join_room(user_id, room_id).await;
        let user_rooms = manager.get_user_rooms(user_id).await;
        assert!(user_rooms.contains(&room_id));

        let room_users = manager.get_room_users(room_id).await;
        assert!(room_users.contains(&user_id));

        // 测试离开房间
        manager.leave_room(user_id, room_id).await;
        let user_rooms = manager.get_user_rooms(user_id).await;
        assert!(!user_rooms.contains(&room_id));

        // 测试移除连接
        manager.remove_connection(user_id).await;
        assert_eq!(manager.get_online_users_count().await, 0);
    }

    #[test]
    fn test_websocket_message_serialization() {
        let message = WebSocketMessage {
            message_type: WebSocketMessageType::Chat,
            content: "Hello, World!".to_string(),
            sender_id: Some(Uuid::new_v4()),
            room_id: Some(Uuid::new_v4()),
            timestamp: chrono::Utc::now(),
        };

        let json = serde_json::to_string(&message).unwrap();
        let deserialized: WebSocketMessage = serde_json::from_str(&json).unwrap();

        assert_eq!(message.content, deserialized.content);
        assert_eq!(message.message_type, deserialized.message_type);
    }
}
